%-------------------------------------------------------------------------%
% This program computes income vectors to feed into NBER taxsim 
%-------------------------------------------------------------------------%

addpath ('...data/HRS MiCDA/Tables');
addpath ('...data');


%Model parameters and grid parameters for W
aT    = 81;            %terminal age
aS    = 50;            %start age
r     = .029;           %real interest rate
pi    = .028;           %inflation rate
ages = aS:1:aT-1;      %ages

%Grid at initial age (aS)
W_gridN = 20;
split = 5;
W_mins = [0;25000;75000;250000;500000;1500000];
Ws=zeros(W_gridN/split,split);
for i = 1:5
    step = (W_mins(i+1)-W_mins(i))/(W_gridN/split);
    Ws(:,i) = W_mins(i):step:(W_mins(i+1)-step);
end
W = reshape(Ws,[W_gridN,1]); 

WGrids = cell(aT-aS+1,1);
WGrids{1,1} = W;
for i = 2:aT-aS+1
    WGrids{i,1} = W.*(1+r)^(i-1);
end

%Differences (pair-wise) to construct contributions and -distributions
dWGrids = cell(aT-aS,1);
dW = zeros(W_gridN,W_gridN);
for a = 1:aT-aS
    for i = 1:W_gridN
        dWGrids{a,1}(i,:) = WGrids{a+1,1}(i)./(1+r) - WGrids{a,1}; %contribution amount
        dWGrids{a,1}(dWGrids{a,1}<0)=0; % zero out negative values
    end
end

%1. HRS db wealth by age
pdata = csvread('compwealth_frz_w10_22.csv',1,0);
%col2 is age
%col3 is avg db wealth in 2010 dollars
%col4 is avg earnings in 2010 dollars
dbw = pdata(aS-48+1:end,2:3);
y_raw = pdata(aS-48+1:end,4);
%Smooth the earnings profile using cubic polynomic fit
y_fit_param =  polyfit(ages',y_raw,3);
y = polyval(y_fit_param,ages');

%2. SS annuity by retirement age (hh-level)
b_ss = csvread('ssw10.csv',1,0);
b_ss = b_ss(aS-48+1:end,4);

%2. SSA mortality rates by age (48+) 
pdeath = csvread('SSA_mortality_2010.csv',1,0);
pdeath = pdeath(aS+1:end,:);
avpdeath = mean(pdeath(:,2:3),2);

%3. annuity factors
afactor = zeros(length(avpdeath),1); 
%loop over age
for a = 1:length(afactor)

    disc = zeros(length(afactor)-a+1,1);
    disc(1) = 1;
    for i = 2:length(disc)
        disc(i) = (1-avpdeath(i))*(1/((1+r)*(1+pi)))^(i-1);
    end

    afactor(a) = sum(disc);
end
afactor = [pdeath(1:aT-aS,1) afactor(1:aT-aS)];

%DB annuity by retirement age
b_db = dbw(:,2)./afactor(:,2);

%-------------------------------------------------------------------------%
% For model inputs without any freeze and average DC plan
%-------------------------------------------------------------------------%

%Pension distribution choices
strt=1;
fin = strt+W_gridN^2-1;
distr = zeros((aT-aS)*W_gridN^2,1);
ss = zeros((aT-aS)*W_gridN^2,1);
age = zeros((aT-aS)*W_gridN^2,1);
for a = 1:length(ages)
    %at each age dWGrids{a,1} + b_db(age) is the set of potential cash outs   
    temp = reshape(b_db(a) + dWGrids{a,1},[],1);
    distr(strt:fin) = temp; 
    %replicate ss ben for all levels of distr per given age
    ss(strt:fin) = b_ss(a);
    %replicate age for all levels of distr
    age(strt:fin) = ages(a);
    clear temp
    strt = fin+1;
    fin = strt+W_gridN^2-1;
end

%DC contribution choices
strt=1;
fin = strt+W_gridN^2-1;
tcontrib = zeros((aT-aS)*W_gridN^2,1);
y_rep = zeros((aT-aS)*W_gridN^2,1);
age = zeros((aT-aS)*W_gridN^2,1);
for a = 1:length(ages)
    %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
    temp = reshape(dWGrids{a,1},[],1);
    tcontrib(strt:fin) = temp;
    %replicate y for all levels of distr per given age
    y_rep(strt:fin) = y(a);
    %total contribution rate
    tcontrib_rate = tcontrib./y_rep;
    %Worker contributions (based on matching function)
    T = (tcontrib_rate<.10135);
    temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
    temp2 = (tcontrib_rate - .03635);
    wcontrib_rate = temp1.*T + temp2.*(1-T);
    wcontrib = max(wcontrib_rate.*y_rep,0);
    %replicate age for all levels of distr
    age(strt:fin) = ages(a);
    clear temp
    strt = fin+1;
    fin = strt+W_gridN^2-1;
end

%Export this set of vectors to .csv to plug into taxsim
csvwrite('.../data/retinc_a_20.csv',[distr ss age]);
csvwrite('.../data/workinc_a_20.csv',[y_rep wcontrib age]);

%-------------------------------------------------------------------------%
% For model inputs when working; different freeze ages and average DC plan
%-------------------------------------------------------------------------%

y_cut = [.037;.020;.015;.048;.016]; 

for aa = 56:64
    
    i_z = find(ages==aa);
    
    %earn profile given freeze at age z
    y_loss = ones(length(y),1);
    y_loss(i_z:i_z+length(y_cut)-1) = 1-y_cut;
    y_z = y.*y_loss;

    %DC contribution choices
    strt=1;
    fin = strt+W_gridN^2-1;
    tcontrib = zeros((aT-aS)*W_gridN^2,1);
    y_rep = zeros((aT-aS)*W_gridN^2,1);
    age = zeros((aT-aS)*W_gridN^2,1);
    for a = 1:length(ages)
        %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
        temp = reshape(dWGrids{a,1},[],1);
        tcontrib(strt:fin) = temp;
        %replicate y for all levels of distr per given age
        y_rep(strt:fin) = y_z(a);
        %total contribution rate
        tcontrib_rate = tcontrib./y_rep;
        %Worker contributions (based on matching function)
        T = (tcontrib_rate<.10135);
        temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
        temp2 = (tcontrib_rate - .03635);
        wcontrib_rate = temp1.*T + temp2.*(1-T);
        wcontrib = max(wcontrib_rate.*y_rep,0);
        %replicate age for all levels of distr
        age(strt:fin) = ages(a);
        clear temp
        strt = fin+1;
        fin = strt+W_gridN^2-1;
    end

    %Export this set of vectors to .csv to plug into taxsim
    fname = strcat('.../data/workinc_a_20_z',num2str(i_z),'.csv');
    csvwrite(fname,[y_rep wcontrib age]);
end

%-------------------------------------------------------------------------%
% For model inputs when working; different freeze ages and TSP replacement
%-------------------------------------------------------------------------%

y_cut = [.037;.020;.015;.048;.016]; 

for aa = 56:64
    
    i_z = find(ages==aa);
    
    %earn profile given freeze at age z
    y_loss = ones(length(y),1);
    y_loss(i_z:i_z+length(y_cut)-1) = 1-y_cut;
    y_z = y.*y_loss;

    %DC contribution choices
    strt=1;
    fin = strt+W_gridN^2-1;
    tcontrib = zeros((aT-aS)*W_gridN^2,1);
    y_rep = zeros((aT-aS)*W_gridN^2,1);
    age = zeros((aT-aS)*W_gridN^2,1);
    
    %Pre freeze
    for a = 1:i_z-1
        %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
        temp = reshape(dWGrids{a,1},[],1);
        tcontrib(strt:fin) = temp;
        %replicate y for all levels of distr per given age
        y_rep(strt:fin) = y_z(a);
        %total contribution rate
        tcontrib_rate = tcontrib./y_rep;
        %Worker contributions (based on matching function)
        T = (tcontrib_rate<.10135);
        temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
        temp2 = (tcontrib_rate - .03635);
        wcontrib_rate = temp1.*T + temp2.*(1-T);
        wcontrib(strt:fin) = max(wcontrib_rate(strt:fin).*y_rep(strt:fin),0);
        %replicate age for all levels of distr
        age(strt:fin) = ages(a);
        clear temp
        strt = fin+1;
        fin = strt+W_gridN^2-1;
    end

    %Post freeze
    for a = i_z:length(ages)
        %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
        temp = reshape(dWGrids{a,1},[],1);
        tcontrib(strt:fin) = temp;
        %replicate y for all levels of distr per given age
        y_rep(strt:fin) = y_z(a);
        %total contribution rate
        tcontrib_rate = tcontrib./y_rep;
        %Worker contributions (based on matching function)
        T1 = (tcontrib_rate<.07);
        T2 = (tcontrib_rate<.1 & tcontrib_rate>=.07);
        temp1 = (tcontrib_rate./(2)) - .01/2;
        temp2 = (tcontrib_rate./(1.5)) - .04/1.5;
        temp3 = (tcontrib_rate - .05);
        wcontrib_rate = temp1.*T1 + temp2.*(1-T2)+temp3.*(1-T1-T2);
        wcontrib(strt:fin) = max(wcontrib_rate(strt:fin).*y_rep(strt:fin),0);
        %replicate age for all levels of distr
        age(strt:fin) = ages(a);
        clear temp
        strt = fin+1;
        fin = strt+W_gridN^2-1;
    end
    
    %Export this set of vectors to .csv to plug into taxsim
    fname = strcat('.../data/workinc_a_20_tspz',num2str(i_z),'.csv');
    csvwrite(fname,[y_rep wcontrib age]);
end


%-------------------------------------------------------------------------%
% For model inputs when working; different freeze ages and new DB contrib
%-------------------------------------------------------------------------%

y_cut = .02; %new 2% contribution


for aa = 56:64
    
    i_z = find(ages==aa);
    
    %earn profile given freeze at age z
    y_loss = ones(length(y),1);
    y_loss(i_z:end) = 1-y_cut;
    y_z = y.*y_loss;

    %DC contribution choices
    strt=1;
    fin = strt+W_gridN^2-1;
    tcontrib = zeros((aT-aS)*W_gridN^2,1);
    y_rep = zeros((aT-aS)*W_gridN^2,1);
    age = zeros((aT-aS)*W_gridN^2,1);
    
    for a = 1:length(ages)
        %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
        temp = reshape(dWGrids{a,1},[],1);
        tcontrib(strt:fin) = temp;
        %replicate y for all levels of distr per given age
        y_rep(strt:fin) = y_z(a);
        %total contribution rate
        tcontrib_rate = tcontrib./y_rep;
        %Worker contributions (based on matching function)
        T = (tcontrib_rate<.10135);
        temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
        temp2 = (tcontrib_rate - .03635);
        wcontrib_rate = temp1.*T + temp2.*(1-T);
        wcontrib(strt:fin) = max(wcontrib_rate(strt:fin).*y_rep(strt:fin),0);
        %replicate age for all levels of distr
        age(strt:fin) = ages(a);
        clear temp
        strt = fin+1;
        fin = strt+W_gridN^2-1;
    end
    
    %Export this set of vectors to .csv to plug into taxsim
    fname = strcat('.../data/workinc_a_20_contz',num2str(i_z),'.csv');
    csvwrite(fname,[y_rep wcontrib age]);
end

%-------------------------------------------------------------------------%
% Retirement income for different freeze ages (56-64)
%-------------------------------------------------------------------------%

for aa = 56:64
    
    i_z = find(ages==aa);
    b_db_z = b_db;
    b_db_z(i_z:end) = b_db(i_z);
    %freeze fixes nominal annuity; compute real annuity at each
    %prospective retirement age
    for i = i_z:length(b_db_z) 
        b_db_z(i) = b_db_z(i)*(1/(1+pi))^(i-i_z);
    end

    %Pension distribution choices
    strt=1;
    fin = strt+W_gridN^2-1;
    distr = zeros((aT-aS)*W_gridN^2,1);
    ss = zeros((aT-aS)*W_gridN^2,1);
    age = zeros((aT-aS)*W_gridN^2,1);
    for a = 1:length(ages)
        %at each age dWGrids{a,1} + b_db(age) is the set of potential cash outs   
        temp = reshape(b_db_z(a) + dWGrids{a,1},[],1);
        distr(strt:fin) = temp; 
        %replicate ss ben for all levels of distr per given age
        ss(strt:fin) = b_ss(a);
        %replicate age for all levels of distr
        age(strt:fin) = ages(a);
        clear temp
        strt = fin+1;
        fin = strt+W_gridN^2-1;
    end

    %Export this set of vectors to .csv to plug into taxsim
    fname = strcat('.../data/retinc_a_20_z',num2str(i_z),'.csv');
    csvwrite(fname,[distr ss age]);
    
end

%-------------------------------------------------------------------------%
%One-time compensation shocks of 20 percent at spec. ages (for elasticities)
%-------------------------------------------------------------------------%

%age 58...
adj = ones(length(ages),1);
adj = [1.*adj(1:8);1.2.*adj(9);1.*adj(10:end)];
y_adj = y.*adj;

%DC contribution choices
strt=1;
fin = strt+W_gridN^2-1;
tcontrib = zeros((aT-aS)*W_gridN^2,1);
y_rep = zeros((aT-aS)*W_gridN^2,1);
age = zeros((aT-aS)*W_gridN^2,1);
for a = 1:length(ages)
    %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
    temp = reshape(dWGrids{a,1},[],1);
    tcontrib(strt:fin) = temp;
    %replicate y for all levels of distr per given age
    y_rep(strt:fin) = y_adj(a);
    %total contribution rate
    tcontrib_rate = tcontrib./y_rep;
    %Worker contributions (based on matching function)
    T = (tcontrib_rate<.10135);
    temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
    temp2 = (tcontrib_rate - .03635);
    wcontrib_rate = temp1.*T + temp2.*(1-T);
    wcontrib = max(wcontrib_rate.*y_rep,0);
    %replicate age for all levels of distr
    age(strt:fin) = ages(a);
    clear temp
    strt = fin+1;
    fin = strt+W_gridN^2-1;
end

%Export this set of vectors to .csv to plug into taxsim
csvwrite('.../data/workinc_s58_a_20.csv',[y_rep wcontrib age]);

%age 60...
adj = ones(length(ages),1);
adj = [1.*adj(1:10);1.2.*adj(11);1.*adj(12:end)];
y_adj = y.*adj;

%DC contribution choices
strt=1;
fin = strt+W_gridN^2-1;
tcontrib = zeros((aT-aS)*W_gridN^2,1);
y_rep = zeros((aT-aS)*W_gridN^2,1);
age = zeros((aT-aS)*W_gridN^2,1);
for a = 1:length(ages)
    %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
    temp = reshape(dWGrids{a,1},[],1);
    tcontrib(strt:fin) = temp;
    %replicate y for all levels of distr per given age
    y_rep(strt:fin) = y_adj(a);
    %total contribution rate
    tcontrib_rate = tcontrib./y_rep;
    %Worker contributions (based on matching function)
    T = (tcontrib_rate<.10135);
    temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
    temp2 = (tcontrib_rate - .03635);
    wcontrib_rate = temp1.*T + temp2.*(1-T);
    wcontrib = max(wcontrib_rate.*y_rep,0);
    %replicate age for all levels of distr
    age(strt:fin) = ages(a);
    clear temp
    strt = fin+1;
    fin = strt+W_gridN^2-1;
end

%Export this set of vectors to .csv to plug into taxsim
csvwrite('.../data/workinc_s60_a_20.csv',[y_rep wcontrib age]);

%age 62...
adj = ones(length(ages),1);
adj = [1.*adj(1:12);1.2.*adj(13);1.*adj(14:end)];
y_adj = y.*adj;

%DC contribution choices
strt=1;
fin = strt+W_gridN^2-1;
tcontrib = zeros((aT-aS)*W_gridN^2,1);
y_rep = zeros((aT-aS)*W_gridN^2,1);
age = zeros((aT-aS)*W_gridN^2,1);
for a = 1:length(ages)
    %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
    temp = reshape(dWGrids{a,1},[],1);
    tcontrib(strt:fin) = temp;
    %replicate y for all levels of distr per given age
    y_rep(strt:fin) = y_adj(a);
    %total contribution rate
    tcontrib_rate = tcontrib./y_rep;
    %Worker contributions (based on matching function)
    T = (tcontrib_rate<.10135);
    temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
    temp2 = (tcontrib_rate - .03635);
    wcontrib_rate = temp1.*T + temp2.*(1-T);
    wcontrib = max(wcontrib_rate.*y_rep,0);
    %replicate age for all levels of distr
    age(strt:fin) = ages(a);
    clear temp
    strt = fin+1;
    fin = strt+W_gridN^2-1;
end

%Export this set of vectors to .csv to plug into taxsim
csvwrite('.../data/workinc_s62_a_20.csv',[y_rep wcontrib age]);

%age 64...
adj = ones(length(ages),1);
adj = [1.*adj(1:14);1.2.*adj(15);1.*adj(16:end)];
y_adj = y.*adj;

%DC contribution choices
strt=1;
fin = strt+W_gridN^2-1;
tcontrib = zeros((aT-aS)*W_gridN^2,1);
y_rep = zeros((aT-aS)*W_gridN^2,1);
age = zeros((aT-aS)*W_gridN^2,1);
for a = 1:length(ages)
    %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
    temp = reshape(dWGrids{a,1},[],1);
    tcontrib(strt:fin) = temp;
    %replicate y for all levels of distr per given age
    y_rep(strt:fin) = y_adj(a);
    %total contribution rate
    tcontrib_rate = tcontrib./y_rep;
    %Worker contributions (based on matching function)
    T = (tcontrib_rate<.10135);
    temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
    temp2 = (tcontrib_rate - .03635);
    wcontrib_rate = temp1.*T + temp2.*(1-T);
    wcontrib = max(wcontrib_rate.*y_rep,0);
    %replicate age for all levels of distr
    age(strt:fin) = ages(a);
    clear temp
    strt = fin+1;
    fin = strt+W_gridN^2-1;
end

%Export this set of vectors to .csv to plug into taxsim
csvwrite('.../data/workinc_s64_a_20.csv',[y_rep wcontrib age]);

%age 66...
adj = ones(length(ages),1);
adj = [1.*adj(1:16);1.2.*adj(17);1.*adj(18:end)];
y_adj = y.*adj;

%DC contribution choices
strt=1;
fin = strt+W_gridN^2-1;
tcontrib = zeros((aT-aS)*W_gridN^2,1);
y_rep = zeros((aT-aS)*W_gridN^2,1);
age = zeros((aT-aS)*W_gridN^2,1);
for a = 1:length(ages)
    %at each age, y - dWGrids{a,1} is the post-DC saving taxable income
    temp = reshape(dWGrids{a,1},[],1);
    tcontrib(strt:fin) = temp;
    %replicate y for all levels of distr per given age
    y_rep(strt:fin) = y_adj(a);
    %total contribution rate
    tcontrib_rate = tcontrib./y_rep;
    %Worker contributions (based on matching function)
    T = (tcontrib_rate<.10135);
    temp1 = (tcontrib_rate./(1.39)) - .011/1.39;
    temp2 = (tcontrib_rate - .03635);
    wcontrib_rate = temp1.*T + temp2.*(1-T);
    wcontrib = max(wcontrib_rate.*y_rep,0);
    %replicate age for all levels of distr
    age(strt:fin) = ages(a);
    clear temp
    strt = fin+1;
    fin = strt+W_gridN^2-1;
end

%Export this set of vectors to .csv to plug into taxsim
csvwrite('.../data/workinc_s66_a_20.csv',[y_rep wcontrib age]);

